package testrank.affinity;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import testrank.TestManager;
import testrank.TestRankConstants;
import testrank.TestRankConstants.AffinityLiteralsSimilarityStrategy;
import testrank.predictors.AffinityContext;


public class AffinityCalculator {

	private static final boolean VERBOSE = false;
	
	private static final double STOPWORDS_FACTOR = 0.25;
	
	private WordNetConnector wnc = new WordNetConnector();
	
	private Map<WordPair, Double> wordAffinities;
	
	private int cacheHits;
	
	private static class GrpAffinity {
		final double totalMaxAffinity;
		final double totalWeight;

		public GrpAffinity(double totalMaxAffinity, double totalWeight) {
			this.totalMaxAffinity = totalMaxAffinity;
			this.totalWeight = totalWeight;
		}

		double get() {
//			if (TestRankConstants.AFFINITY_ALT_NORMALIZATION == TestRankConstants.AffinityNormalization.Recursive) {
//				return totalMaxAffinity;
//			}
			
			double asyGroupAffinity = totalWeight != 0 ? totalMaxAffinity  / totalWeight : 0.0;
			return asyGroupAffinity;
		}
	}
	
	public AffinityCalculator() {
		TestManager.loadAffinities(); 
		wordAffinities = TestManager.wordAffinities();
		if (wordAffinities == null) {
			wordAffinities = new HashMap<WordPair, Double>();
			TestManager.setWordAffinities(wordAffinities);
		}
	}
	
	public double wordAffinityIgnoreCase(String w1, String w2, Set<String> g1, Set<String> g2) {
		return wordAffinity(w1.toLowerCase(), w2.toLowerCase(), g1, g2);
	}
	
	public double wordAffinity(String w1, String w2, Set<String> g1, Set<String> g2) {
		WordPair key = new WordPair(w1, w2);
		//CAHCE
		Double wordAffinity = wordAffinities.get(key);
		if (wordAffinity == null) {
			int dist = wnc.wordDistance(w1, w2, g1, g2);
			wordAffinity = 1.0 / dist; 
				//dist != WordNetConnector.DISTANCE_NOT_FOUND ? 1.0 / dist : 0.0;
			// CAHCE
			wordAffinities.put(key, wordAffinity);
			if (cacheHits > 0) {
				if (VERBOSE) {
					System.out.println("cacheHits="+cacheHits);
				}
				cacheHits = 0;
			}
			if (VERBOSE) {
				System.out.println("wordAffinity "+key+" -> " + wordAffinity);
			}
		}
		else {		
			cacheHits++;
		}
		
		return wordAffinity;
	}
	
	public double groupAffinity(Map<String, Double> g1, Map<String, Double> g2, WordGroup wg1, WordGroup wg2, AffinityContext affinityContext, boolean useWordAffinity, boolean verbose) {
		if (TestRankConstants.GRP_AFFINITY_CALC == TestRankConstants.GroupAffinityCalc.Cosine) {
			return cosineGroupAffinity(g1, g2, wg1, wg2, affinityContext, useWordAffinity, verbose);
		}
		
		GrpAffinity asyGrpAff1 = asyGroupAffinity(g1, g2, wg1, wg2, affinityContext, useWordAffinity, verbose); 
		GrpAffinity asyGrpAff2 = asyGroupAffinity(g2, g1, wg2, wg1, affinityContext, useWordAffinity, verbose);
		double groupAffinity;
		if (TestRankConstants.GRP_AFFINITY_CALC == 
				TestRankConstants.GroupAffinityCalc.AsyAvg1 ||
			TestRankConstants.GRP_AFFINITY_CALC == 
				TestRankConstants.GroupAffinityCalc.AsyAvg3) {
			groupAffinity = (asyGrpAff1.get() + asyGrpAff2.get()) / 2;
		}
		else if (TestRankConstants.GRP_AFFINITY_CALC == 
			TestRankConstants.GroupAffinityCalc.AsyAvg2 ||
				 TestRankConstants.GRP_AFFINITY_CALC == 
			TestRankConstants.GroupAffinityCalc.AsyAvg4) {
			GrpAffinity a = new GrpAffinity(asyGrpAff1.totalMaxAffinity + asyGrpAff2.totalMaxAffinity, 
					asyGrpAff1.totalWeight + asyGrpAff2.totalWeight);
			groupAffinity = a.get();
			assert 0.0 <= groupAffinity && groupAffinity <= 1.0; 
		}
		else throw new IllegalStateException();
		
		return groupAffinity;
	}
	
	private GrpAffinity asyGroupAffinity(Map<String, Double> g1, Map<String, Double> g2, WordGroup wg1, WordGroup wg2, AffinityContext affinityContext, boolean useWordAffinity, boolean verbose) {
		if (g1.size() == 0 || g2.size() == 0) {
			return new GrpAffinity(0.0, 0.0);
		}

//		g1 = WordNetConnector.reduceWordsWithSameBaseFrom(g1);
//		g2 = WordNetConnector.reduceWordsWithSameBaseFrom(g2);
		
		double totalMaxAffinity = 0.0;
		double totalWeight = 0.0;
		double squareSum1 = 0.0;
		double squareSum2 = 0.0;
		for (String w1 : g1.keySet()) {
	
			double tfidf1 = useWordAffinity ? 
					wg1.getWordTFIDF(w1, affinityContext.allCoveredMethodsOfTestWordGroups) : 
					wg1.getLiteralTFIDF(w1, affinityContext.allCoveredMethodsOfTestWordGroups);
			
			if (useWordAffinity) {
				tfidf1 = fixTfidfIfStopword(w1, tfidf1);
			}
			
			double maxWordAffinity = -1.0;
			String maxAffinityW2 = null;
			
			double maxTfidf2 = 0.0;
		
			for (String w2 : g2.keySet()) {
				
				double tfidf2 = useWordAffinity ? 
					wg2.getWordTFIDF(w2, affinityContext.allCoveringTestsOfMethodWordGroups) : 
					wg2.getLiteralTFIDF(w2, affinityContext.allCoveringTestsOfMethodWordGroups);
				
				double wordAffinity = useWordAffinity ?
					wordAffinityIgnoreCase(w1, w2, g1.keySet(), g2.keySet()):
					literalsAffinity(w1, w2);
					
					
				if (verbose) {
					System.out.println("a(" + w1 + ", " + w2 + ") = " + wordAffinity);
				}
				
				if (wordAffinity > maxWordAffinity) {
					maxWordAffinity = wordAffinity;
					maxAffinityW2 = w2;

					maxTfidf2 = tfidf2;
				}
			}

			if (useWordAffinity) {
				maxTfidf2 = fixTfidfIfStopword(maxAffinityW2, maxTfidf2);
			}
			
			// g1.get(w1) is method name weights 
			double weight = tfidf1*maxTfidf2*g1.get(w1)*g2.get(maxAffinityW2);
			totalMaxAffinity += maxWordAffinity * weight;
			totalWeight += weight;
			
			double weight1 = tfidf1*g1.get(w1);
			double weight2 = maxTfidf2*g2.get(maxAffinityW2);
			squareSum1 += weight1*weight1;
			squareSum2 += weight2*weight2;
		}
		
		
		if (TestRankConstants.GRP_AFFINITY_CALC == 
				TestRankConstants.GroupAffinityCalc.AsyAvg3 ||
			TestRankConstants.GRP_AFFINITY_CALC == 
				TestRankConstants.GroupAffinityCalc.AsyAvg4) {
			totalWeight = Math.sqrt(squareSum1 * squareSum2); 
		}

		GrpAffinity asyGroupAffinity = new GrpAffinity(totalMaxAffinity, totalWeight);	
		
		return asyGroupAffinity;
	}

	private static class TermAffinity {
		String term;
		double affinity;

		public TermAffinity(String term, double affinity) {
			this.term = term;
			this.affinity = affinity;
		}
		
	}
	private double cosineGroupAffinity(Map<String, Double> g1,
			Map<String, Double> g2, WordGroup wg1, WordGroup wg2,
			AffinityContext affinityContext, boolean useWordAffinity,
			boolean verbose) {
		if (g1.size() == 0 || g2.size() == 0) {
			return 0.0;
		}
		
		SortedSet<String> terms = new TreeSet<String>();
		terms.addAll(g1.keySet());
		terms.addAll(g2.keySet());
		
		double[][] weights = new double[2][terms.size()];
		fillWeights(g1, wg1, affinityContext, useWordAffinity, verbose, terms, weights[0]);
		fillWeights(g2, wg2, affinityContext, useWordAffinity, verbose, terms, weights[1]);
		
		normalizeWeights(weights[0]);
		normalizeWeights(weights[1]);
		
		double dotProduct = 0.0;
		for (int i = 0; i < weights[0].length; i++) {
			dotProduct += weights[0][i] * weights[1][i];
		}

		assert 0.0 <= dotProduct && dotProduct <= 1.0 : ("illegal dot product: " + dotProduct);
		System.out.println(dotProduct);
		return dotProduct;
	}


	private double normalizeWeights(double[] v) {
		double squareSum = 0.0;
		for (double x : v) {
			squareSum += x*x;
		}
		double norm = Math.sqrt(squareSum);
		for (int i = 0; i < v.length; i++) {
			if (norm > 0.0) {
				v[i] /= norm; 
			}
			else {
				v[i] = 0.0;
			}
		}
			
		return norm;
	}

	private void fillWeights(Map<String, Double> g, WordGroup wg,
			AffinityContext affinityContext, boolean useWordAffinity,
			boolean verbose, SortedSet<String> terms, double[] weights) {
		int i = 0;
		for (String term : terms) {
			double aff1; // affinity of term from g1 closet to term in g2
			if (g.containsKey(term)) {
				aff1 = 1.0;
			}
			else { // term appears only in g2: find the most similar term instead
				TermAffinity termAffinity = getMaxAffinityTerm(g, term, useWordAffinity, verbose);
				term = termAffinity.term;
				aff1 = termAffinity.affinity;
			}
			double termWeight = getTermWeight(term, g, wg, affinityContext, useWordAffinity);
			weights[i] = aff1 * termWeight;
			
			i++;
		}
	}

	
	private TermAffinity getMaxAffinityTerm(Map<String, Double> g, String w, boolean useWordAffinity, boolean verbose) {
		double maxAffinity = -1.0;
		String maxAffinityTerm = null;
		
	
		for (String w2 : g.keySet()) {
			if (w2 == w) {
				continue;
			}
			
			double wordAffinity = useWordAffinity ?
				wordAffinityIgnoreCase(w, w2, null, null):
				literalsAffinity(w, w2);
				
				
			if (verbose) {
				System.out.println("a(" + w + ", " + w2 + ") = " + wordAffinity);
			}
			
			if (wordAffinity > maxAffinity) {
				maxAffinity = wordAffinity;
				maxAffinityTerm = w2;
			}
		}
		
		assert maxAffinity != -1;
		
		return new TermAffinity(maxAffinityTerm, maxAffinity);
	}

	private double getTermWeight(String term, Map<String, Double> g, WordGroup wg,
			AffinityContext affinityContext, boolean useWordAffinity) {
		double tfidf = useWordAffinity ? 
				wg.getWordTFIDF(term, affinityContext.allCoveredMethodsOfTestWordGroups) : 
				wg.getLiteralTFIDF(term, affinityContext.allCoveredMethodsOfTestWordGroups);
		
		if (useWordAffinity) {
			tfidf = fixTfidfIfStopword(term, tfidf);
		}
		
		return tfidf * g.get(term);
	}

	// TODO somehow these words get here, despite the stopwords filter in SourceFileAnalyzer:
	// get, to, do, is
	private double fixTfidfIfStopword(String w, double tfidf) {
		if (/*w.length() > 1 &&*/ SourceFileAnalyzer.isStopword(w)) {
			//System.out.println("STOPWORD: " + w);
			tfidf *= STOPWORDS_FACTOR; 
			// isn't really different than tfidf=0,
			// since currently stopwords are not stored in the first place...  
		}
		return tfidf;
	}


	private double literalsAffinity(String w1, String w2) {
		if (TestRankConstants.AFFINITY_LITERALS_SIMILARITY_STRATEGY == 
			AffinityLiteralsSimilarityStrategy.EqualsIgnoreCase) {
			return w1.equalsIgnoreCase(w2) ? 1.0 : 0.0;	
		}
		else if (TestRankConstants.AFFINITY_LITERALS_SIMILARITY_STRATEGY == 
			AffinityLiteralsSimilarityStrategy.EditDistance) {
			return EditDistanceCalculator.getSimilarity(w1, w2, false);
		}
		else throw new IllegalStateException();

//		if (w1.equals(w2))
//		return 1.0; 
//	else return w1.equalsIgnoreCase(w2) ? 0.5 : 0.0;
	}
	
	public static void main(String[] args) {
		AffinityCalculator ac = new AffinityCalculator(); 
		measureGroupAffinity(ac, "testing", "bike", "speed", null, "vehicle", "velocity");
		measureGroupAffinity(ac, "hello", "world", null, "hello", "goodbye");
		measureGroupAffinity(ac, "hello", "world", null, "expanding", "universe");
		
		resetProfiling();
		
		measureGroupAffinity(ac, "expand", "wildcard", "field", null, "expanding", "star", "attribute");
		measureGroupAffinity(ac, "square", "root", null, "square", "equation");
		measureGroupAffinity(ac, "negative", "fail", null, "validate");
		measureGroupAffinity(ac, "select", "auto", "top", "down", null, "dig", "in", "selection", "listener");
		measureGroupAffinity(ac, "unselect", "child", "negative", null, "dig", "in", "remove", "selection");
		measureGroupAffinity(ac, "connection", null, "open", "close", "session");
		measureGroupAffinity(ac, "long", "time", "ago", null, "i", "cry", "music");
		measureGroupAffinity(ac, "knew", "chance", null, "make", "people", "dance");
		measureGroupAffinity(ac, "maybe", "happy", null, "february", "shiver", "paper", "deliver");
		measureGroupAffinity(ac, "news", "doorstep", null, "make", "one", "more", "step");
		measureGroupAffinity(ac, "bye", "miss", "american", "pie", null, "dry", "good", "old", "boys");
		measureGroupAffinity(ac, "write", "book", null, "love", "faith", "god");
		measureGroupAffinity(ac, "above", "bible", "tell", null, "believe", "rock", "roll");
		measureGroupAffinity(ac, "music", "save", "mortal", null, "soul", "teach", "dance", "real", "slow");
		measureGroupAffinity(ac, "know", "love", "him", "saw", null, "dancing", "gym");
		measureGroupAffinity(ac, "both", "kicked", "shoes", null, "rythm", "blues");
	}
	private static void resetProfiling() {
		
	}

	private static void measureGroupAffinity(AffinityCalculator ac, String... words) {
		
		Set<String> g1 = new HashSet<String>();
		Set<String> g2 = new HashSet<String>();
		Set<String> s = g1;
		for (String w : words) {
			if (w == null)
				s = g2;
			else
				s.add(w); 
		}
		//System.out.println("Affinity between " + g1 + " and " + g2 + " is " + ac.groupAffinity(g1, g2, null, null, true, false));
		System.out.println();
	}
}
